#!/usr/bin/env python3
"""
Extract additional mapping relationships from partially decrypted text
"""

import os
import json
import re
from collections import Counter, defaultdict

def analyze_partial_decryption():
    """Analyze partial decryption to find more mapping relationships"""
    
    print("=== Analyzing Partial Decryption Results ===")
    
    # Load original encrypted data
    with open('decoded_result.bin', 'rb') as f:
        encrypted_data = f.read()
    
    # Load partial decryption result
    with open('partial_mapping_candidate_1_0x2fba5b0.txt', 'rb') as f:
        partial_result = f.read()
    
    print(f"Original encrypted size: {len(encrypted_data)} bytes")
    print(f"Partial result size: {len(partial_result)} bytes")
    
    # Known mappings we already have
    known_mappings = {
        0xa1: 0x6d,  # m
        0xf3: 0x6f,  # o  
        0x6b: 0x64,  # d
        0x31: 0x65,  # e
        0xb0: 0x6c,  # l
        0x69: 0x22,  # "
        0x01: 0x3a,  # :
        0x90: 0x20   # space
    }
    
    # Analyze patterns in the partial result
    analyze_json_patterns(partial_result, encrypted_data, known_mappings)
    
    # Find common character patterns
    find_character_patterns(partial_result, encrypted_data, known_mappings)
    
    # Analyze repeated sequences
    analyze_repeated_sequences(partial_result, encrypted_data, known_mappings)

def analyze_json_patterns(partial_result, encrypted_data, known_mappings):
    """Look for JSON patterns to infer more mappings"""
    
    print("\n=== JSON Pattern Analysis ===")
    
    try:
        # Convert to string for pattern analysis
        result_str = partial_result.decode('utf-8', errors='ignore')
        
        # Look for JSON-like patterns
        json_patterns = [
            ('"model":', 'Field name with colon'),
            (',"', 'Field separator'),
            ('"}', 'Object closing'),
            ('{"', 'Object opening'),
            ('":"', 'Key-value separator'),
            ('":', 'Field colon'),
        ]
        
        for pattern, desc in json_patterns:
            positions = []
            start = 0
            while True:
                pos = result_str.find(pattern, start)
                if pos == -1:
                    break
                positions.append(pos)
                start = pos + 1
                if len(positions) > 20:  # Limit results
                    break
            
            if positions:
                print(f"\nPattern '{pattern}' ({desc}): {len(positions)} occurrences")
                
                # Analyze the encrypted bytes at these positions
                for pos in positions[:5]:  # First 5 occurrences
                    if pos < len(encrypted_data) - len(pattern):
                        encrypted_seq = encrypted_data[pos:pos+len(pattern)]
                        result_seq = result_str[pos:pos+len(pattern)]
                        print(f"  Position {pos}: {encrypted_seq.hex()} -> '{result_seq}'")
                        
                        # Extract individual character mappings
                        extract_mappings_from_sequence(encrypted_seq, result_seq.encode('utf-8'), known_mappings)
    
    except Exception as e:
        print(f"Error in JSON pattern analysis: {e}")

def extract_mappings_from_sequence(encrypted_seq, expected_seq, known_mappings):
    """Extract mappings from aligned sequences"""
    
    new_mappings = {}
    
    for i in range(min(len(encrypted_seq), len(expected_seq))):
        enc_byte = encrypted_seq[i]
        exp_byte = expected_seq[i]
        
        if enc_byte not in known_mappings:
            new_mappings[enc_byte] = exp_byte
    
    if new_mappings:
        print(f"    New mappings: {[(hex(k), chr(v) if 32 <= v <= 126 else hex(v)) for k, v in new_mappings.items()]}")
    
    return new_mappings

def find_character_patterns(partial_result, encrypted_data, known_mappings):
    """Find common character patterns to infer mappings"""
    
    print("\n=== Character Pattern Analysis ===")
    
    try:
        result_str = partial_result.decode('utf-8', errors='ignore')
        
        # Count character frequencies in partial result
        char_freq = Counter(result_str)
        common_chars = char_freq.most_common(20)
        
        print("Most frequent characters in partial result:")
        for char, freq in common_chars:
            if 32 <= ord(char) <= 126:  # Printable ASCII
                print(f"  '{char}': {freq} times")
        
        # Look for English word patterns
        print("\nLooking for English word patterns...")
        word_pattern = re.compile(r'\b[a-zA-Z]{2,}\b')
        words = word_pattern.findall(result_str)
        word_freq = Counter(words)
        
        print("Potential English words found:")
        for word, freq in word_freq.most_common(10):
            if len(word) > 2:
                print(f"  '{word}': {freq} times")
                
                # Find positions of this word
                positions = []
                start = 0
                while True:
                    pos = result_str.find(word, start)
                    if pos == -1:
                        break
                    positions.append(pos)
                    start = pos + 1
                    if len(positions) > 3:
                        break
                
                # Analyze encrypted bytes for this word
                for pos in positions[:2]:
                    if pos < len(encrypted_data) - len(word):
                        encrypted_word = encrypted_data[pos:pos+len(word)]
                        print(f"    '{word}' at {pos}: {encrypted_word.hex()}")
                        
                        # Extract new mappings
                        extract_mappings_from_sequence(encrypted_word, word.encode('utf-8'), known_mappings)
    
    except Exception as e:
        print(f"Error in character pattern analysis: {e}")

def analyze_repeated_sequences(partial_result, encrypted_data, known_mappings):
    """Analyze repeated sequences to find more mappings"""
    
    print("\n=== Repeated Sequence Analysis ===")
    
    try:
        # Look for repeated 3-byte sequences in encrypted data
        seq_positions = defaultdict(list)
        
        for i in range(len(encrypted_data) - 2):
            seq = encrypted_data[i:i+3]
            seq_positions[seq].append(i)
        
        # Find sequences that appear multiple times
        repeated_seqs = {seq: positions for seq, positions in seq_positions.items() 
                        if len(positions) > 1}
        
        print(f"Found {len(repeated_seqs)} repeated 3-byte sequences")
        
        # Analyze the most common repeated sequences
        sorted_seqs = sorted(repeated_seqs.items(), key=lambda x: len(x[1]), reverse=True)
        
        for seq, positions in sorted_seqs[:10]:  # Top 10
            print(f"\nSequence {seq.hex()} appears {len(positions)} times")
            
            # Check what it decodes to in different positions
            decoded_variants = set()
            for pos in positions[:5]:  # Check first 5 occurrences
                if pos + 3 <= len(partial_result):
                    decoded_seq = partial_result[pos:pos+3]
                    try:
                        decoded_str = decoded_seq.decode('utf-8', errors='ignore')
                        if all(32 <= ord(c) <= 126 for c in decoded_str):
                            decoded_variants.add(decoded_str)
                    except:
                        pass
            
            if decoded_variants:
                print(f"  Decodes to: {decoded_variants}")
                
                # If it consistently decodes to the same thing, extract mappings
                if len(decoded_variants) == 1:
                    decoded = list(decoded_variants)[0]
                    if len(decoded) == 3:
                        print(f"  Consistent mapping: {seq.hex()} -> '{decoded}'")
                        extract_mappings_from_sequence(seq, decoded.encode('utf-8'), known_mappings)
    
    except Exception as e:
        print(f"Error in repeated sequence analysis: {e}")

def create_improved_mapping_table():
    """Create improved mapping table with all discovered mappings"""
    
    print("\n=== Creating Improved Mapping Table ===")
    
    # Start with known mappings
    improved_mappings = {
        0xa1: 0x6d,  # m
        0xf3: 0x6f,  # o  
        0x6b: 0x64,  # d
        0x31: 0x65,  # e
        0xb0: 0x6c,  # l
        0x69: 0x22,  # "
        0x01: 0x3a,  # :
        0x90: 0x20   # space
    }
    
    # Add likely mappings based on frequency analysis
    # These are educated guesses based on common JSON characters
    likely_mappings = {
        # Common JSON characters
        # We need to find the encrypted bytes for these
    }
    
    # For now, let's try to find mappings for common characters
    # by analyzing the partial decryption
    
    with open('decoded_result.bin', 'rb') as f:
        encrypted_data = f.read()
    
    with open('partial_mapping_candidate_1_0x2fba5b0.txt', 'rb') as f:
        partial_result = f.read()
    
    # Look for specific patterns we know should exist
    search_patterns = [
        (b'"model"', '6d6f64656c'),  # "model"
        (b'":"', '3a'),              # ":"  
        (b',"', '2c'),               # ","
        (b'}', '7d'),                # }
        (b'{', '7b'),                # {
    ]
    
    for pattern_bytes, hex_desc in search_patterns:
        # Find this pattern in the partial result
        pos = partial_result.find(pattern_bytes)
        if pos != -1 and pos < len(encrypted_data) - len(pattern_bytes):
            encrypted_seq = encrypted_data[pos:pos+len(pattern_bytes)]
            print(f"Found pattern '{pattern_bytes.decode()}' at position {pos}")
            print(f"  Encrypted as: {encrypted_seq.hex()}")
            
            # Add to improved mappings
            for i, enc_byte in enumerate(encrypted_seq):
                if i < len(pattern_bytes):
                    expected_byte = pattern_bytes[i]
                    if enc_byte not in improved_mappings:
                        improved_mappings[enc_byte] = expected_byte
                        print(f"  New mapping: 0x{enc_byte:02x} -> '{chr(expected_byte)}' (0x{expected_byte:02x})")
    
    # Create full 256-byte mapping table
    full_table = bytearray(range(256))  # Start with identity
    
    for src, dst in improved_mappings.items():
        full_table[src] = dst
    
    # Save improved mapping table
    with open('improved_mapping.bin', 'wb') as f:
        f.write(full_table)
    
    print(f"\nCreated improved mapping table with {len(improved_mappings)} known mappings")
    print("Saved to improved_mapping.bin")
    
    # Test the improved mapping
    test_improved_mapping(bytes(full_table))
    
    return improved_mappings

def test_improved_mapping(mapping_table):
    """Test the improved mapping table"""
    
    print("\n=== Testing Improved Mapping ===")
    
    with open('decoded_result.bin', 'rb') as f:
        encrypted_data = f.read()
    
    try:
        # Apply the improved mapping
        decrypted = bytes(mapping_table[b] for b in encrypted_data)
        
        # Check if we can find JSON patterns
        decrypted_str = decrypted.decode('utf-8', errors='ignore')
        
        # Look for improvements
        json_indicators = ['"model"', '":', ',"', '{"', '"}']
        found_patterns = sum(1 for pattern in json_indicators if pattern in decrypted_str)
        
        print(f"Found {found_patterns}/{len(json_indicators)} JSON patterns")
        
        if found_patterns >= 3:
            # Save the result
            with open('improved_decryption.txt', 'w', encoding='utf-8', errors='ignore') as f:
                f.write(decrypted_str)
            
            print("✅ Improved decryption saved to improved_decryption.txt")
            
            # Show preview
            preview = decrypted_str[:500]
            print(f"Preview: {preview[:200]}...")
            
            return True
        else:
            print("❌ No significant improvement found")
            
    except Exception as e:
        print(f"Error testing improved mapping: {e}")
    
    return False

if __name__ == "__main__":
    print("Extracting Additional Mapping Relationships")
    print("=" * 50)
    
    if not os.path.exists('decoded_result.bin'):
        print("❌ decoded_result.bin not found")
        exit(1)
    
    if not os.path.exists('partial_mapping_candidate_1_0x2fba5b0.txt'):
        print("❌ Partial decryption result not found")
        exit(1)
    
    # Analyze the partial decryption
    analyze_partial_decryption()
    
    # Create improved mapping table
    improved_mappings = create_improved_mapping_table()
    
    print(f"\n{'=' * 50}")
    print("Analysis complete!")
    print(f"Total mappings discovered: {len(improved_mappings)}")
    print("Check improved_decryption.txt for results")